import numpy as np
from config import xcfg as cfg
import numpy as np
from matplotlib import pyplot as plt
from scipy.sparse.linalg import spsolve
from scipy import sparse
import pypardiso
import torch
import time
if cfg.method_params.solver == 'numpy_gpu':
    import cupy as cp

def metric(fvm_results, exact_results):
    error = np.mean(np.abs(fvm_results - exact_results))
    error /= np.mean(fvm_results)
    return error


def problem_solve(cfg):
    N = cfg.method_params.num_control_uints
    L = cfg.L
    TA = cfg.TA
    TB = cfg.TB

    k = np.full([N], cfg.k)
    q = np.full([N], cfg.q) 

    dx = L / N

    # 线性方程组的右边项初始化
    _b = np.zeros(N) 

    # 系数矩阵
    coeffic_matrix = np.zeros([N, N], dtype = np.float32)

    # 离散系数整理
    for i in range(1, N-1):
        a_W = -(k[i-1] + k[i]) / 2 / dx
        a_E = -(k[i] + k[i+1]) / 2 / dx
        b_C = q[i] * dx

        # 三个系数
        coeffic_matrix[i, i-1] = a_W
        coeffic_matrix[i, i+1] = a_E
        coeffic_matrix[i, i] = -(a_W + a_E)

        # 右边项
        _b[i] = b_C
    
    # A边界的处理
    coeffic_matrix[0, 1] = -(k[0] + k[1]) / 2 / dx
    _b[0] = 2 * k[0] / dx * TA + q[0] * dx
    coeffic_matrix[0, 0] = -coeffic_matrix[0, 1] + 2 * k[0] / dx

    # B边界的处理
    coeffic_matrix[N-1, N-2] = -(k[N-2] + k[N-1]) / 2 / dx
    _b[N-1] = k[N-1] * 2 / dx * TB + q[N-1] * dx
    coeffic_matrix[N-1, N-1] = -coeffic_matrix[N-1, N-2] + k[N-1] * 2 / dx

    print('Start to solve the linear equation. ==>')
    st_time = time.time()
    if cfg.method_params.solver == 'numpy':
        result = np.linalg.solve(coeffic_matrix, _b) # 解线性方程组 AX = b
    elif cfg.method_params.solver == 'sparse_scipy':
        sp_coeffic_matrix = sparse.csr_matrix(coeffic_matrix)
        result = spsolve(sp_coeffic_matrix, _b)
    elif cfg.method_params.solver == 'sparse_pardiso':
        sp_coeffic_matrix = sparse.csr_matrix(coeffic_matrix)
        coeffic_matrix = coeffic_matrix.astype(np.float64)
        _b = _b.astype(np.float64)
        result = pypardiso.spsolve(sp_coeffic_matrix, _b)
    elif cfg.method_params.solver == 'numpy_gpu':
        """ GPU memory >= 19GB """
        coeffic_matrix = cp.asarray(coeffic_matrix)
        _b = cp.asarray(_b)
        result = cp.linalg.solve(coeffic_matrix, _b)
        result = cp.asnumpy(result)

    elif cfg.method_params.solver == 'torch_gpu':
        """ GPU memory >= 19GB """
        coeffic_matrix = torch.from_numpy(coeffic_matrix).cuda()
        _b = torch.from_numpy(_b).cuda()
        result = torch.linalg.solve(coeffic_matrix, _b)
        result = result.cpu().numpy()

    ed_time = time.time()
    cost_time = ed_time - st_time
    if cfg.report.solver_time:
        print('solve over. ==> using time: {}'.format(cost_time))

    x = np.linspace(0, L, N + 2)
    result = np.array([TA, *result, TB])
    exact = ((TB-TA)/L+q[0]/2/k[0]*(L-x)) * x + TA
    error = metric(result, exact)
    if cfg.report.error_print:
        print("Error: {}".format(error))
    if cfg.report.fvm_fig:
        plt.plot(x, result, 'o', label='FVM')
        if cfg.report.exact_fig:
            plt.plot(x, exact, '-', label='Exact')
            plt.legend()
            plt.show()

if __name__ == '__main__':
    problem_solve(cfg)



